(*pp camlp4orf *)
open Camlp4
open PreCast
module A = PreCast.Ast
module P = Pa_type_conv
module Gen = P.Gen

let apply_unit _loc x = 
  <:expr< $x$ () >>
;;
let abstract_unit _loc x = 
  <:expr< fun () -> $x$ >>
;;
let expr_of_list  _loc exprs = 
  let rec coll =
    function
	[] ->
	  <:expr< [] >>
      | x :: xs ->
	  <:expr< [ $x$ :: $coll xs$ ] >>
  in
    coll exprs
;;

let fold1 f lst  = 
  let rec loop sofar =
    function
      | [] ->
	  sofar
      | x :: xs ->
	  loop (f sofar x) xs
  in
    match lst with
      | [] -> 
	  failwith "fold1 failed"
      | x :: xs ->
	  loop x xs

;;

let expr_of_tuple _loc exprs = 
  <:expr<$tup:Ast.exCom_of_list exprs$>>
;;

module Specification = 
struct

  let mk_abst_call _loc tn rev_path =
    <:expr< $id:Gen.ident_of_rev_path _loc (("specification_of_" ^ tn) :: rev_path)$ >>
  ;;
      (* Conversion of type paths *)
      let of_path_fun _loc id =
	match Gen.get_rev_id_path id [] with
	  | ["unit"] -> <:expr< Test.Conv.Specification.of_unit >>
	  | ["bool"] -> <:expr< Test.Conv.Specification.of_bool >>
	  | ["string"] -> <:expr< Test.Conv.Specification.of_string >>
	  | ["list"] -> <:expr< Test.Conv.Specification.of_list >>
	  | ["char"] -> <:expr< Test.Conv.Specification.of_char >>
	  | ["int"] -> <:expr< Test.Conv.Specification.of_int >>
	  | ["float"] -> <:expr< Test.Conv.Specification.of_float >>
	  | tn :: rev_path -> mk_abst_call _loc tn rev_path
	  | [] -> assert false  (* impossible *)

  (* Conversion of types *)
  let rec of_type  _loc = 
    function
	(*    | <:ctyp< sexp_opaque $_$ >> ->  <:expr< Sexplib.Conv.sexp_of_opaque >>*)
      | <:ctyp< $tp1$ $tp2$ >> ->  (of_appl_fun  _loc tp1 tp2)
      | <:ctyp< ( $tup:tp$ ) >> -> of_tuple  _loc tp
      | <:ctyp< '$parm$ >> ->  <:expr< $lid:("_of_" ^ parm)$ >>
      | <:ctyp< $id:id$ >> ->  (of_path_fun _loc id)
      | <:ctyp< $tp1$ -> $tp2$ >> ->  failwith "not supported" 
      | <:ctyp< [< $row_fields$ ] >> | <:ctyp< [> $row_fields$ ] >>
      | <:ctyp< [= $row_fields$ ] >> -> failwith "eh"
      | <:ctyp< ! $parms$ . $poly_tp$ >> -> failwith "don't even get this"
      | _ ->
          prerr_endline (P.get_loc_err _loc "of_type: unknown type construct in of_type");
          exit 1

  and  of_sum _loc  alts   =
    let constructors = 
      let rec fold _loc = 
	function
	  | <:ctyp< $uid:cnstr$ | $tp2$ >> ->
	    (cnstr,[]) :: (fold _loc tp2)
	  | <:ctyp< $uid:cnstr$ of $tp$  | $tp2$ >> ->
	    (cnstr,A.list_of_ctyp tp []) :: (fold _loc tp2)
	  | <:ctyp< $uid:cnstr$  >> ->
	    [(cnstr,[])]
	  | <:ctyp< $uid:cnstr$ of $tp$ >> ->
	    [(cnstr,A.list_of_ctyp tp [])]
	  | _ -> failwith "branch_sum: unknown type"
      in
	fold _loc alts 
    in 
    let constructors = 
      List.map
	(fun (cnstr,tps) ->
	   <:expr<
	     ($str:cnstr$,
	      $expr_of_list _loc
		(List.map (of_type _loc) tps)$)
	   >>)
	constructors
    in
      abstract_unit _loc
      <:expr< Test.Specification.Variant $expr_of_list _loc constructors$ >>
	


  (* Conversion of polymorphic types *)
  and of_appl_fun  _loc tp1 tp2 =
    match of_type  _loc tp1, of_type  _loc tp2 with
      | fun_expr1, fun_expr2 -> <:expr< $fun_expr1$ $fun_expr2$ >>

  (* Conversion of tuples *)
	  
  and of_tuple  _loc tp = 
    let tps = A.list_of_ctyp tp [] in
    let tps = List.map (of_type _loc) tps in 
      abstract_unit _loc
      <:expr< Test.Specification.Tuple $expr_of_list _loc tps$ >>

  and of_record _loc flds_ctyp = 
    let flds = A.list_of_ctyp flds_ctyp [] in
    let flds = 
      let rec coll sofar = function
	| [] -> sofar
	| <:ctyp< $lid:name$ : mutable $tp$ >> :: xs
	| <:ctyp< $lid:name$ : $tp$ >> :: xs ->
	  coll ((name, tp)  :: sofar) xs
	| _ -> raise (Failure ("could not convert record"))
      in
	List.rev (coll [] flds)
    in
    let flds = 
      List.map 
	(fun (str,tp) -> <:expr< ($str:str$,$of_type _loc tp$) >>) 
	flds
    in
      abstract_unit _loc
      <:expr< Test.Specification.Record $expr_of_list _loc flds$ >>      
  ;;	
  (* Empty type *)
  let of_nil _loc =  <:expr< fun _v -> assert False >>
  ;;
  let of_td  _loc type_name tps rhs =
    let is_alias_ref = ref false in
    let handle_alias _loc tp = is_alias_ref := true; of_type  _loc tp in

    let body = 
      let rec loop _loc =
	Gen.switch_tp_def _loc
	  ~alias:handle_alias
	  ~sum:of_sum
	  ~record:of_record
	  ~variants:(fun _ -> assert false)
	  ~mani:(fun _loc _tp1 -> loop _loc)
	  ~nil:of_nil
      in
	loop _loc rhs
    in
    let patts =
      List.map 
	(function
	     <:ctyp< '$x$ >> ->
	   <:patt< ($lid:"_of_" ^ x$) >>
	   | _ -> raise (Failure "weird"))
	tps
    in
    let body = <:expr< fun x -> $body$ x >> in
    let bnd = Gen.idp _loc ("specification_of_" ^ type_name) in
      <:binding< $bnd$ =  $Gen.abstract _loc patts body$ >>
  ;;	
  let rec of_tds  = function
    | A.TyDcl (_loc, type_name, tps, rhs, _cl) ->
	of_td  _loc type_name tps rhs
    | A.TyAnd (_loc, tp1, tp2) ->
	<:binding< $of_tds  tp1$ and $of_tds  tp2$ >>
	| _ -> assert false  (* impossible *)

end

module Representation = 
struct
  let mk_abst_call _loc tn rev_path =
    <:expr< $id:Gen.ident_of_rev_path _loc ((tn ^ "_of_representation") :: rev_path)$ >>


      let rec make_vars i = 
	if i= 0 then [] else ("x" ^ (string_of_int i)) :: make_vars (i-1) 
  ;;

  let patt_of_list _loc lst = 
    let rec loop  = 
      function
	| [] -> <:patt< [] >>
	| x :: xs ->
	    <:patt< [ $x$ :: $loop xs$ ] >>
    in
      loop lst 
  ;;
  let make_list_pattern _loc i = 
    let vars = make_vars i in 
    let evars = 
      List.map 
	(fun v -> <:patt< $lid:v$ >>)
	vars
    in
      patt_of_list _loc evars,vars
  ;;	

  let expr_of_record _loc lst = 
    let rec iter =
      function
	| [] -> raise (Failure "no such thing as an empty record")
	| [(x,y)] ->
	    <:rec_binding< $lid:x$ = $y$ >>
	| (x,y) :: xs ->
	    let z = <:rec_binding< $lid:x$ = $y$ >> in
	      <:rec_binding< $z$ ; $iter xs$  >>
    in
      <:expr< { $iter lst$ } >>


	(* Conversion of type paths *)
	let of_path_fun _loc id =
	  match Gen.get_rev_id_path id [] with
	    | ["unit"] -> <:expr< Test.Conv.Representation.to_unit >>
	    | ["bool"] -> <:expr< Test.Conv.Representation.to_bool >>
	    | ["string"] -> <:expr< Test.Conv.Representation.to_string >>
	    | ["list"] -> <:expr< Test.Conv.Representation.to_list >>
	    | ["char"] -> <:expr< Test.Conv.Representation.to_char >>
	    | ["int"] -> <:expr< Test.Conv.Representation.to_int >>
	    | ["float"] -> <:expr< Test.Conv.Representation.to_float >>
	    | tn :: rev_path -> mk_abst_call _loc tn rev_path
	    | [] -> assert false  (* impossible *)

  (* Conversion of types *)
  let rec of_type  _loc = 
    function
	(*    | <:ctyp< sexp_opaque $_$ >> ->  <:expr< Sexplib.Conv.sexp_of_opaque >>*)
      | <:ctyp< $tp1$ $tp2$ >> ->  (of_appl_fun  _loc tp1 tp2)
      | <:ctyp< ( $tup:tp$ ) >> -> of_tuple  _loc tp
      | <:ctyp< '$parm$ >> ->  <:expr< $lid:("_of_" ^ parm)$ >>
      | <:ctyp< $id:id$ >> ->  (of_path_fun _loc id)
      | <:ctyp< $tp1$ -> $tp2$ >> ->  failwith "not supported" 
      | <:ctyp< [< $row_fields$ ] >> | <:ctyp< [> $row_fields$ ] >>
      | <:ctyp< [= $row_fields$ ] >> -> failwith "eh"
      | <:ctyp< ! $parms$ . $poly_tp$ >> -> failwith "don't even get this"
      | _ ->
          prerr_endline (P.get_loc_err _loc "of_type: unknown type construct in of_type");
          exit 1

  and  of_sum _loc alts =
    let constructors = 
      let rec fold _loc = 
	function
	  | <:ctyp< $uid:cnstr$ | $tp2$ >> ->
	    (cnstr,[]) :: (fold _loc tp2)
	  | <:ctyp< $uid:cnstr$ of $tp$  | $tp2$ >> ->
	    (cnstr,A.list_of_ctyp tp []) :: (fold _loc tp2)
	  | <:ctyp< $uid:cnstr$  >> ->
	    [(cnstr,[])]
	  | <:ctyp< $uid:cnstr$ of $tp$ >> ->
	    [(cnstr,A.list_of_ctyp tp [])]
	  | _ -> failwith "branch_sum: unknown type"
      in
	fold _loc alts 
    in 
    let cases = 
      let f (cnstr,tps) = 
	match tps with
	  | [] -> <:match_case<Test.Representation.Variant ($str:cnstr$,[]) -> $uid:cnstr$>>
	  | _ :: _ ->
	let patt,vars = make_list_pattern _loc (List.length tps) in 
	let tps = List.combine tps vars in
	let tps = List.map (fun (t,v) -> <:expr<$of_type _loc t$ $lid:v$>>) tps in 
	let body = <:expr<$uid:cnstr$ $tup:Ast.exCom_of_list tps$ >> in 
	  <:match_case<Test.Representation.Variant ($str:cnstr$,$patt$) -> $body$ >>
      in
	List.map
	  f
	  constructors
    in
    let base_case =
	<:match_case< _ -> raise (Test.Of_representation "variant") >> 
    in
    let match_body = 
      fold1 
	(fun a b -> <:match_case< $a$ | $b$ >>) 
	(cases @ [base_case])  
    in
      <:expr<
	fun x ->
	  match x with [$match_body$] >>
	      
  (* Conversion of tuples *)
	      
  and of_tuple  _loc tp = 
    let tps = A.list_of_ctyp tp [] in
    let tps = List.map (of_type _loc) tps in 
    let pattern, variables = make_list_pattern _loc (List.length tps) in  
    let tps = List.combine variables tps in 
    let match_body = 
      List.map
	(fun (v,t) ->
	   <:expr<$t$ $lid:v$>>)
	tps
    in
      <:expr< 
	fun x ->
	  match x with
	      [ Test.Representation.Tuple($pattern$)  -> $tup:Ast.exCom_of_list match_body$
	      | _ -> raise (Test.Of_representation "tuple" )
	      ]
	      >>

		
  and of_record _loc flds_ctyp = 
    let make_record_pattern lst = 
      let vars = make_vars (List.length lst) in 
      let vars = List.combine lst vars in 
      let triplets = 
	List.map
	  (fun ((cnstr,t),v) ->
	     let patt = 
	       <:patt< ($str:cnstr$,$lid:v$) >>
	     in
	     let expr = 
	       <:expr< $of_type _loc t$ $lid:v$ >>
	     in
	       (patt,(cnstr,expr)))
	  vars
      in
      let patts = List.map fst triplets in 
      let body = List.map snd triplets in 
	patt_of_list _loc patts, body
    in


    let flds = A.list_of_ctyp flds_ctyp [] in
    let flds = 
      let rec coll sofar = function
	| [] -> sofar
	| <:ctyp< $lid:name$ : mutable $tp$ >> :: xs
	| <:ctyp< $lid:name$ : $tp$ >> :: xs ->
	  coll ((name, tp)  :: sofar) xs
	| _ -> raise (Failure ("could not convert record"))
      in
	List.rev (coll [] flds)
    in
    let patt,record = make_record_pattern flds in 
    let record = expr_of_record _loc record in 
      <:expr< 
	fun x ->
	  match x with
	      [ Test.Representation.Record $patt$  -> $record$
	      | _ -> raise (Test.Of_representation "record" )
	      ]
	      >>

  (* Conversion of polymorphic types *)
  and of_appl_fun  _loc tp1 tp2 =
    match of_type  _loc tp1, of_type  _loc tp2 with
      | fun_expr1, fun_expr2 -> <:expr< fun x -> $fun_expr1$ $fun_expr2$ x >>
  ;;
  (* Empty type *)
  let of_nil _loc =  <:expr< fun _v -> assert False >>
  ;;
  let of_td  _loc type_name tps rhs =
    let is_alias_ref = ref false in
    let handle_alias _loc tp = is_alias_ref := true; of_type  _loc tp in

    let body = 
      let rec loop _loc =
	Gen.switch_tp_def _loc
	  ~alias:handle_alias
	  ~sum:of_sum
	  ~record:of_record
	  ~variants:(fun _ -> assert false)
	  ~mani:(fun _loc _tp1 -> loop _loc)
	  ~nil:of_nil
      in
	loop _loc rhs
    in
    let patts =
      List.map 
	(function
	     <:ctyp< '$x$ >> ->
	   <:patt< ($lid:"_of_" ^ x$) >>
	   | _ -> raise (Failure "weird"))
	tps
    in

    let bnd = Gen.idp _loc (type_name ^ "_of_representation"  ) in
      <:binding< $bnd$ =  $Gen.abstract _loc patts body$ >>
  ;;
  let rec of_tds  = function
    | A.TyDcl (_loc, type_name, tps, rhs, _cl) ->
	of_td  _loc type_name tps rhs
    | A.TyAnd (_loc, tp1, tp2) ->
	<:binding< $of_tds  tp1$ and $of_tds  tp2$ >>
	| _ -> assert false  (* impossible *)
  ;;
end

module Generate_specification_of = struct
  let specification_of  tds =
    let binding, recursive, _loc =
      match tds with
	| A.TyDcl (_loc, type_name, tps, rhs, _cl) ->
            Specification.of_td  _loc type_name tps rhs,
            Gen.type_is_recursive _loc type_name rhs, _loc
	| A.TyAnd (_loc, _, _) as tds -> Specification.of_tds  tds, true, _loc
	| _ -> assert false  (* impossible *)
    in
      <:str_item< value rec $binding$ >>
  ;;
  let string_of_ident id =
    let str_lst = Gen.get_rev_id_path id [] in
      String.concat "." str_lst
end

module Generate_representation_of = struct
  let representation_of  tds =
    let binding, recursive, _loc =
      match tds with
	| A.TyDcl (_loc, type_name, tps, rhs, _cl) ->
            Representation.of_td  _loc type_name tps rhs,
            Gen.type_is_recursive _loc type_name rhs, _loc
	| A.TyAnd (_loc, _, _) as tds -> Representation.of_tds  tds, true, _loc
	| _ -> assert false  (* impossible *)
    in
      <:str_item< value rec $binding$ >>
  ;;
  
  let string_of_ident id =
    let str_lst = Gen.get_rev_id_path id [] in
      String.concat "." str_lst
end

let () =
  P.add_generator
    "tui"
    (fun tds ->
       let _loc = Loc.ghost in
	 <:str_item<
	   $Generate_specification_of.specification_of  tds$;
	 $Generate_representation_of.representation_of  tds$
	 >>)
    
